
import argparse
import time
from typing import List, Dict, Optional
from datetime import datetime, timedelta
import pandas as pd
import logging
from qopilot.data.database import SQLiteManagerHK, TableSchema
from qopilot.data.datasource.qop_akshare import HKAKShareData
from qopilot.constant import TradeDate, Adjust
from qopilot.utils import write_to_excel
"""
港股数据收集脚本

Author: Ryan
Email: 409266255@qq.com
Date: 2025-01-28 10:00:00

Description:
    用于初始化港股数据库，以及增量更新港股列表和个股历史行情数据
    个股历史行情数据表的命名规则为 {stock_code}.HK_{adjust}

Features:
    1. 数据库初始化
    2. 港股列表获取和更新
    3. 个股历史行情数据获取
    4. 增量更新逻辑
    5. 错误处理和重试机制
"""

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class HKDataCollector:
    """港股数据收集器"""

    def __init__(self):
        """初始化数据收集器"""
        self.db_manager = SQLiteManagerHK()
        self.data_source = HKAKShareData()
        self.retry_count = 3
        self.retry_delay = 5  # 秒

        logger.info("港股数据收集器初始化完成")

    def init_database(self) -> bool:
        """
        初始化港股数据库

        Returns:
            bool: 初始化是否成功
        """
        try:
            logger.info("开始初始化港股数据库")

            # 初始化港股列表表
            self.db_manager.create_table(
                table_name='hk_stocks_list',
                columns=TableSchema.HK_STOCKS_LIST_SCHEMA.value,
                primary_key='code',
                unique_constraints=['code']
            )

            logger.info("港股数据库初始化完成")
            return True

        except Exception as e:
            logger.error(f"数据库初始化失败: {str(e)}")
            return False

    def create_stock_hist_table(self, stock_code: str, adjust: str = Adjust.QFQ.value) -> bool:
        """
        为指定股票代码创建历史行情数据表

        Args:
            stock_code: 股票代码

        Returns:
            bool: 创建是否成功
        """
        try:
            table_name = f"{adjust}_{stock_code}_HK"

            # 检查表是否已存在
            if self.db_manager.table_exists(table_name):
                logger.info(f"表 {table_name} 已存在")
                return True

            # 创建表
            self.db_manager.create_table(
                table_name=table_name,
                columns=TableSchema.HK_STOCK_HIST_SCHEMA.value,
                primary_key='date',
                unique_constraints=['date']
            )

            logger.info(f"成功创建表: {table_name}")
            return True

        except Exception as e:
            logger.error(f"创建表失败 {stock_code}: {str(e)}")
            return False

    def get_hk_stocks_list(self, force_update: bool = False) -> Optional[pd.DataFrame]:
        """
        获取港股列表

        Args:
            force_update: 是否强制更新

        Returns:
            DataFrame: 港股列表数据
        """
        try:
            logger.info("开始获取港股列表")

            # 检查是否需要更新
            if not force_update and self._is_stocks_list_up_to_date():
                logger.info("港股列表数据已是最新，跳过更新")
                return self._get_stocks_list_from_db()

            # 获取最新港股列表
            df_stocks = self.data_source.get_hk_stocks_list()

            if df_stocks is not None and not df_stocks.empty:
                # 保存到数据库
                self._save_stocks_list_to_db(df_stocks)
                logger.info(f"成功获取港股列表，共 {len(df_stocks)} 只股票")
                return df_stocks
            else:
                logger.warning("获取港股列表失败，返回空数据")
                return None

        except Exception as e:
            logger.error(f"获取港股列表失败: {str(e)}")
            return None

    def _is_stocks_list_up_to_date(self) -> bool:
        """检查港股列表是否是最新的（今天的数据）"""
        try:
            today = datetime.now().strftime('%Y%m%d')
            sql = "SELECT COUNT(*) as count FROM hk_stocks_list WHERE date = ?"
            result = self.db_manager.query_data(sql, (today,))

            if result and result[0]['count'] > 0:
                return True
            return False

        except Exception as e:
            logger.error(f"检查港股列表更新状态失败: {str(e)}")
            return False

    def _get_stocks_list_from_db(self) -> Optional[pd.DataFrame]:
        """从数据库获取港股列表"""
        try:
            sql = "SELECT * FROM hk_stocks_list ORDER BY code"
            result = self.db_manager.query_data(sql)

            if result:
                df = pd.DataFrame(result)
                logger.info(f"从数据库获取港股列表，共 {len(df)} 只股票")
                return df
            return None

        except Exception as e:
            logger.error(f"从数据库获取港股列表失败: {str(e)}")
            return None

    def _save_stocks_list_to_db(self, df: pd.DataFrame):
        """保存港股列表到数据库"""
        try:
            self.db_manager.insert_dataframe(
                df=df,
                table_name='hk_stocks_list',
                if_exists='replace',
                index=False
            )
            logger.info(f"成功保存港股列表到数据库，共 {len(df)} 只股票")
        except Exception as e:
            logger.error(f"保存港股列表到数据库失败: {str(e)}")
            excel_filename = f"hk_stocks_name_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
            write_to_excel(df, excel_filename)
            raise

    def get_stock_hist_data(self, stock_code: str,
                            start_date: str = TradeDate.START_DATE.value,
                            end_date: str = None,
                            force_update: bool = False,
                            adjust: str = Adjust.QFQ.value) -> Optional[pd.DataFrame]:
        """
        获取单只股票的历史行情数据

        Args:
            stock_code: 股票代码
            start_date: 开始日期
            end_date: 结束日期，默认为今天
            force_update: 是否强制更新

        Returns:
            DataFrame: 历史行情数据
        """
        try:
            if end_date is None:
                end_date = datetime.now().strftime('%Y%m%d')

            logger.info(f"开始获取股票 {stock_code} 的历史行情数据")

            # 创建表（如果不存在）
            if not self.create_stock_hist_table(stock_code, adjust):
                return None

            # 检查是否需要更新
            if not force_update and self._is_stock_hist_up_to_date(stock_code, end_date, adjust):
                logger.info(f"股票 {stock_code} 历史数据已是最新，跳过更新")
                return self._get_stock_hist_from_db(stock_code, adjust)

            # 获取历史数据
            df_hist = self.data_source.get_hk_hist(
                code=stock_code,
                start_date=start_date,
                end_date=end_date,
                adjust=adjust
            )

            if df_hist is not None and not df_hist.empty:
                # 保存到数据库
                self._save_stock_hist_to_db(stock_code, df_hist, adjust)
                logger.info(f"成功获取股票 {stock_code} 历史数据，共 {len(df_hist)} 条记录")
                return df_hist
            else:
                logger.warning(f"获取股票 {stock_code} 历史数据失败，返回空数据")
                return None

        except Exception as e:
            logger.error(f"获取股票 {stock_code} 历史数据失败: {str(e)}")
            return None

    def _is_stock_hist_up_to_date(self, stock_code: str, end_date: str, adjust: str = Adjust.QFQ.value) -> bool:
        """检查股票历史数据是否是最新的"""
        try:
            table_name = f"{adjust}_{stock_code}_HK"
            sql = f"SELECT MAX(date) as max_date FROM {table_name}"
            result = self.db_manager.query_data(sql)

            if result and result[0]['max_date']:
                max_date = result[0]['max_date']
                return max_date >= end_date
            return False

        except Exception as e:
            logger.error(f"检查股票 {stock_code} 历史数据更新状态失败: {str(e)}")
            return False

    def _get_stock_hist_from_db(self, stock_code: str, adjust: str = Adjust.QFQ.value) -> Optional[pd.DataFrame]:
        """从数据库获取股票历史数据"""
        try:
            table_name = f"{adjust}_{stock_code}_HK"
            sql = f"SELECT * FROM {table_name} ORDER BY date"
            result = self.db_manager.query_data(sql)

            if result:
                df = pd.DataFrame(result)
                logger.info(f"从数据库获取股票 {stock_code} 历史数据，共 {len(df)} 条记录")
                return df
            return None

        except Exception as e:
            logger.error(f"从数据库获取股票 {stock_code} 历史数据失败: {str(e)}")
            return None

    def _save_stock_hist_to_db(self, stock_code: str, df: pd.DataFrame, adjust: str = Adjust.QFQ.value):
        """保存股票历史数据到数据库"""
        try:
            table_name = f"{adjust}_{stock_code}_HK"

            # 添加创建时间
            df['created_at'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            # 保存到数据库
            self.db_manager.insert_dataframe(
                df=df,
                table_name=table_name,
                if_exists='replace',
                index=False
            )

            logger.info(f"✅ 成功保存股票 {stock_code} 历史数据到数据库")

        except Exception as e:
            logger.error(f"❌ 保存股票 {stock_code} 历史数据到数据库失败: {str(e)}")
            raise

    def batch_update_stocks_hist(self, stock_codes: List[str] = None,
                                 start_date: str = TradeDate.START_DATE.value,
                                 end_date: str = None,
                                 force_update: bool = False,
                                 adjust: str = Adjust.QFQ.value) -> Dict[str, bool]:
        """
        批量更新股票历史数据

        Args:
            stock_codes: 股票代码列表，如果为None则更新所有股票
            start_date: 开始日期
            end_date: 结束日期
            force_update: 是否强制更新

        Returns:
            Dict[str, bool]: 更新结果字典
        """
        try:
            if stock_codes is None:
                # 获取所有股票代码
                df_stocks = self._get_stocks_list_from_db()
                if df_stocks is None or df_stocks.empty:
                    logger.error("无法获取股票列表")
                    return {}
                stock_codes = df_stocks['code'].tolist()

            logger.info(f"开始批量更新 {len(stock_codes)} 只股票的历史数据")

            results = {}
            success_count = 0

            for i, stock_code in enumerate(stock_codes, 1):
                try:
                    logger.info(f"正在更新第 {i}/{len(stock_codes)} 只股票: {stock_code}")

                    df_hist = self.get_stock_hist_data(
                        stock_code=stock_code,
                        start_date=start_date,
                        end_date=end_date,
                        force_update=force_update,
                        adjust=adjust
                    )

                    if df_hist is not None and not df_hist.empty:
                        results[stock_code] = True
                        success_count += 1
                        logger.info(f"股票 {stock_code} 更新成功")
                    else:
                        results[stock_code] = False
                        logger.warning(f"股票 {stock_code} 更新失败")

                    # 添加延迟避免接口限流
                    if i < len(stock_codes):
                        logger.info("⚠️ 等待7秒")
                        time.sleep(7)

                except Exception as e:
                    logger.error(f"更新股票 {stock_code} 失败: {str(e)}")
                    results[stock_code] = False

            logger.info(f"批量更新完成，成功 {success_count}/{len(stock_codes)} 只股票")
            return results

        except Exception as e:
            logger.error(f"批量更新股票历史数据失败: {str(e)}")
            return {}

    def incremental_update(self, days_back: int = 1) -> bool:
        """
        增量更新最近几天的数据

        Args:
            days_back: 回溯天数

        Returns:
            bool: 更新是否成功
        """
        try:
            logger.info(f"开始增量更新最近 {days_back} 天的数据")

            # 计算日期范围
            end_date = datetime.now()
            start_date = end_date - timedelta(days=days_back)

            start_date_str = start_date.strftime('%Y%m%d')
            end_date_str = end_date.strftime('%Y%m%d')

            logger.info(f"增量更新日期范围: {start_date_str} - {end_date_str}")

            # 获取股票列表
            df_stocks = self._get_stocks_list_from_db()
            if df_stocks is None or df_stocks.empty:
                logger.error("无法获取股票列表，增量更新失败")
                return False

            # 批量更新
            results = self.batch_update_stocks_hist(
                stock_codes=df_stocks['code'].tolist(),
                start_date=start_date_str,
                end_date=end_date_str,
                force_update=True
            )

            success_count = sum(1 for success in results.values() if success)
            total_count = len(results)

            logger.info(f"增量更新完成，成功 {success_count}/{total_count} 只股票")
            return success_count > 0

        except Exception as e:
            logger.error(f"增量更新失败: {str(e)}")
            return False

    def get_database_stats(self) -> Dict[str, any]:
        """获取数据库统计信息"""
        try:
            stats = {}

            # 港股列表统计
            sql = "SELECT COUNT(*) as count FROM hk_stocks_list"
            result = self.db_manager.query_data(sql)
            stats['stocks_count'] = result[0]['count'] if result else 0

            # 获取所有历史数据表
            sql = "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE '%.HK_%'"
            tables = self.db_manager.query_data(sql)
            stats['hist_tables_count'] = len(tables)

            # 计算总记录数
            total_records = 0
            for table in tables:
                table_name = table['name']
                sql = f"SELECT COUNT(*) as count FROM {table_name}"
                result = self.db_manager.query_data(sql)
                if result:
                    total_records += result[0]['count']

            stats['total_hist_records'] = total_records

            return stats

        except Exception as e:
            logger.error(f"获取数据库统计信息失败: {str(e)}")
            return {}


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='港股数据收集脚本')
    parser.add_argument('--init', action='store_true', help='初始化数据库')
    parser.add_argument('--update-stocks', action='store_true', help='更新港股列表')
    parser.add_argument('--update-hist', nargs='?', const='all',
                        help='更新历史数据，可指定股票代码或使用all更新所有')
    parser.add_argument('--incremental', type=int, default=1,
                        help='增量更新最近N天的数据')
    parser.add_argument('--force', action='store_true', help='强制更新')
    parser.add_argument('--stats', action='store_true', help='显示数据库统计信息')

    args = parser.parse_args()

    # 创建数据收集器
    collector = HKDataCollector()

    try:
        # 初始化数据库
        if args.init:
            if collector.init_database():
                print("✅ 数据库初始化成功")
                return
            else:
                print("❌ 数据库初始化失败")
                return

        # 更新港股列表
        if args.update_stocks:
            df_stocks = collector.get_hk_stocks_list(force_update=args.force)
            if df_stocks is not None:
                print(f"✅ 港股列表更新成功，共 {len(df_stocks)} 只股票")
                return
            else:
                print("❌ 港股列表更新失败")
                return

        # 更新历史数据
        if args.update_hist:
            if args.update_hist == 'all':
                # TODO: 面对限流，需要增加重试机制
                results = collector.batch_update_stocks_hist(force_update=args.force, adjust=Adjust.QFQ.value)
                success_count = sum(1 for success in results.values() if success)
                print(f"✅ 历史数据更新完成，成功 {success_count}/{len(results)} 只股票")
                return
            else:
                df_hist = collector.get_stock_hist_data(
                    stock_code=args.update_hist,
                    force_update=args.force,
                    adjust=Adjust.QFQ.value
                )
                if df_hist is not None:
                    print(f"✅ 股票 {args.update_hist} 历史数据更新成功，共 {len(df_hist)} 条记录")
                    return
                else:
                    print(f"❌ 股票 {args.update_hist} 历史数据更新失败")
                    return

        # 增量更新
        if args.incremental:
            if collector.incremental_update(days_back=args.incremental):
                print(f"✅ 增量更新成功，更新最近 {args.incremental} 天的数据")
            else:
                print("❌ 增量更新失败")

        # 显示统计信息
        if args.stats:
            stats = collector.get_database_stats()
            print("\n📊 数据库统计信息:")
            print(f"  港股数量: {stats.get('stocks_count', 0)}")
            print(f"  历史数据表数量: {stats.get('hist_tables_count', 0)}")
            print(f"  历史数据总记录数: {stats.get('total_hist_records', 0)}")

        # 如果没有指定任何操作，显示帮助信息
        if not any([args.init, args.update_stocks, args.update_hist, args.incremental, args.stats]):
            parser.print_help()

    except KeyboardInterrupt:
        print("\n⚠️  用户中断操作")
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
        print(f"❌ 程序执行失败: {str(e)}")


if __name__ == "__main__":
    main()
